import os
from flask import Flask, render_template, request, jsonify
from sparkai.llm.llm import ChatSparkLLM, ChunkPrintHandler
from sparkai.core.messages import ChatMessage

#星火认知大模型 SPARKAI_DOMAIN 和 SPARKAI_URL 映射关系，请参考文档（https://www.xfyun.cn/doc/spark/Web.html）
SPARKAI_URL_CONFIG = {
    "general": "wss://spark-api.xf-yun.com/v1.1/chat",
    "generalv2": "wss://spark-api.xf-yun.com/v2.1/chat",
    "generalv3": "wss://spark-api.xf-yun.com/v3.1/chat",
    "generalv3.5": "wss://spark-api.xf-yun.com/v3.5/chat",
    "pro-128": "wss://spark-api.xf-yun.com/chat/pro-128k",
    "4.0Ultra": "wss://spark-api.xf-yun.com/v4.0/chat"
}

#星火认知大模型Spark Max的domain值，其他版本大模型domain值请前往文档（https://www.xfyun.cn/doc/spark/Web.html）查看
SPARKAI_DOMAIN = os.getenv('ext_cf_spark_llm_domain')
SPARKAI_URL = SPARKAI_URL_CONFIG.get(SPARKAI_DOMAIN)
#星火认知大模型调用秘钥信息，请前往讯飞开放平台控制台（https://console.xfyun.cn/services/bm35）查看
SPARKAI_APP_ID = os.getenv('ext_cf_spark_app_id')
SPARKAI_API_SECRET = os.getenv('ext_cf_spark_api_secret')
SPARKAI_API_KEY = os.getenv('ext_cf_spark_api_key')

app = Flask(__name__)

@app.route('/')
def index():
    return render_template('chat.html')

@app.route('/send_message', methods=['POST'])
def handle_message():
    spark = ChatSparkLLM(
        spark_api_url=SPARKAI_URL,
        spark_app_id=SPARKAI_APP_ID,
        spark_api_key=SPARKAI_API_KEY,
        spark_api_secret=SPARKAI_API_SECRET,
        spark_llm_domain=SPARKAI_DOMAIN,
        streaming=False,
    )

    data = request.json  # 假设客户端发送的是 JSON 格式的数据
    print('Message: ' + data['message'])
    # 这里可以添加逻辑来处理消息，例如存储到数据库等
    messages = [ChatMessage(
        role="user",
        content=data['message']
    )]
    handler = ChunkPrintHandler()
    retMsg = spark.generate([messages], callbacks=[handler])
    return jsonify({'status': 'success', 'message': retMsg.generations[0][0].text}), 200


if __name__ == '__main__':
    app.run(host='0.0.0.0', port=8000, debug=True)

